from DDPM.Config import *
import numpy as np
import DEAD.AutoDecoder.Evaluation.InternalBallistic as IB
from Utility.PlotCurve import plot_curves
import re
import os
from DDPM.Evaluation.Control import *
import torch

def convert_av1_to_value(av1):
    """
    将av1向量转换为实际的r_ref和R值。
    
    参数:
    - av1: 包含r_ref和R的归一化向量
    
    返回:
    - r_ref: 推进剂燃烧（燃速）特性
    - R: 装药半径
    """
    av1 = av1 * 10.0
    r_ref = (r_ref_max - r_ref_min) * av1[0] + r_ref_min
    R = (R_max - R_min) * av1[1] + R_min
    return r_ref, R

def convert_av2_to_value(av2):
    """
    将av2向量转换为实际的cstar_rho, n和At值。
    
    参数:
    - av2: 包含cstar_rho, n和At的归一化向量
    
    返回:
    - cstar_rho: 特征速度乘以密度
    - n: 压力指数
    - At: 喷口面积
    """
    cstar_rho = (cstar_max * rho_max - cstar_min * rho_min) * av2[0] + cstar_min * rho_min
    n = (n_max - n_min) * av2[1] + n_min
    At = av2[2] * 0.01
    return cstar_rho, n, At

def get_av2_values_from_index(trainingset_path, psfd_index, item_index):
    """
    从索引获取av2值。
    
    参数:
    - trainingset_path: 训练集路径
    - psfd_index: PSFD索引
    - item_index: 项目索引
    
    返回:
    - cstar_rho: 特征速度乘以密度
    - n: 压力指数
    - At: 喷口面积
    """
    npz_path = f'{trainingset_path}/{psfd_index}'
    npz_file = np.load(npz_path + ".npz")
    av2 = npz_file["avs2"][item_index]
    cstar_rho, n, At = convert_av2_to_value(av2)
    return cstar_rho, n, At

def get_av1_values_from_generated_samples(working_path, sample_index):
    """
    从生成的样本中获取av1值。
    
    参数:
    - working_path: 工作路径
    - sample_index: 样本索引
    
    返回:
    - r_ref: 推进剂燃烧（燃速）特性
    - R: 装药半径
    """
    try:
        npz_file = np.load(f"{working_path}/Samples_clustered.npz")
    except FileNotFoundError:
        try:
            npz_file = np.load(f"{working_path}/Samples.npz")
        except FileNotFoundError:
            npz_file = np.load(f"{working_path}/Optimization.npz")
    av1 = npz_file["av1"][sample_index]
    r_ref, R = convert_av1_to_value(av1)
    return r_ref, R

def get_target_from_generated_samples(working_path):
    """
    从生成的样本中获取目标内弹道曲线。
    
    参数:
    - working_path: 工作路径
    
    返回:
    - t_array_target: 目标时间数组
    - p_array_target: 目标压力数组
    """
    try:
        npz_file = np.load(f"{working_path}/Samples_clustered.npz")
    except FileNotFoundError:
        try:
            npz_file = np.load(f"{working_path}/Samples.npz")
        except FileNotFoundError:
            npz_file = np.load(f"{working_path}/Optimization.npz")
    p_array_target = npz_file["p_array_target"]
    t_array_target = npz_file["t_array_target"]
    return t_array_target, p_array_target

def get_av2_values_from_generated_samples(working_path):
    """
    从生成的样本中获取av2值。
    
    参数:
    - working_path: 工作路径
    
    返回:
    - cstar: 特征速度
    - rho: 密度
    - n: 压力指数
    - At: 喷口面积
    """
    try:
        npz_file = np.load(f"{working_path}/Samples_clustered.npz")
    except FileNotFoundError:
        try:
            npz_file = np.load(f"{working_path}/Samples.npz")
        except FileNotFoundError:
            npz_file = np.load(f"{working_path}/Optimization.npz")
    cstar = npz_file["cstar"]
    rho = npz_file["rho"]
    n = npz_file["n"]
    At = npz_file["At"]
    return cstar, rho, n, At

def ballistic_curve_from_index(trainingset_path, psfd_index, item_index):
    """
    从索引获取弹道曲线。
    
    参数:
    - trainingset_path: 训练集路径
    - psfd_index: PSFD索引
    - item_index: 项目索引
    
    返回:
    - t_array: 时间数组
    - p_array: 压力数组
    """
    npz_path = f'{trainingset_path}/{psfd_index}'
    npz_file = np.load(npz_path + ".npz")
    bv = npz_file["bvs"][item_index]
    t_array = bv[:, 0]
    p_array = bv[:, 1] * 1e6
    return t_array, p_array

def start_internal_ballistic_cal(control_params):
    print(f"##### Start internal ballistic cal of {control_params['project_name']} #####")
    project_name=control_params["project_name"]
    user_defined=control_params["user_defined"]
    from OPT.Objective import performance_matching_degree
    project_path=f"{decode_result_save_path}/{project_name}"
    if user_defined:
        t_array_target, p_array_target = get_target_from_generated_samples(project_path)
        cstar, rho, n, At = get_av2_values_from_generated_samples(project_path)
        cstar_rho = cstar * rho
    else:
        i_from_psfd=control_params["i_from_psfd"]
        cstar_rho, n, At = get_av2_values_from_index(trainingset_path, i_from_psfd, 0)
        t_array_target, p_array_target = ballistic_curve_from_index(trainingset_path, i_from_psfd, 0)

    # 获取目录中的所有文件
    files = os.listdir(f"{project_path}")
    # 过滤出 .npz 结尾的文件
    npz_files = [file for file in files if file.endswith('.npz') and re.search('#', file)]
    npz_files = sorted(npz_files, key=lambda x: int(re.search(r'(\d+)#', x).group(1)))

    print(npz_files)

    t_array_list = []
    p_array_list = []
    r2_list=[]
    r2_k_list=[]
    for i in range(len(npz_files)):
        index_in_samples = int(re.search(r'(\d+)#', npz_files[i]).group(1))
        r_ref, R = get_av1_values_from_generated_samples(project_path, index_in_samples)
        npz_file = np.load(f"{project_path}/{npz_files[i]}")
        w_array_norm = npz_file['burned_web_array']
        Ab_array_norm = npz_file['burning_surface_area_array']
        t_array_designed, p_array_designed = IB.internal_ballistic_calculation(
            w_array_norm=w_array_norm,
            Ab_array_norm=Ab_array_norm,
            cstar_rho=cstar_rho,
            r_ref=r_ref,
            n=n,
            R=R,
            At=At,
            verify_range=False)
        t_array_list.append(t_array_designed)
        p_array_list.append(p_array_designed)

        t_array_target_tensor = torch.tensor(t_array_target, dtype=torch.float32)
        p_array_target_tensor = torch.tensor(p_array_target, dtype=torch.float32)
        t_array_designed_tensor = torch.tensor(t_array_designed, dtype=torch.float32)
        p_array_designed_tensor = torch.tensor(p_array_designed, dtype=torch.float32)

        with torch.no_grad():
            r2, r2_k, t_vector_designed, p_vector_designed=\
                performance_matching_degree(t_array_designed_tensor,
                                        p_array_designed_tensor,
                                        t_array_target_tensor,
                                        p_array_target_tensor)
            r2_list.append(r2.item())
            r2_k_list.append(r2_k.item())
    np.savez_compressed(f"{project_path}/internal_ballistic.npz",
                    t_array_design=t_array_list,
                    p_array_design=p_array_list,
                    t_array_target=t_array_target,
                    p_array_target=p_array_target,
                    r2_k_list=r2_k_list,
                    r2_list=r2_list,
                    r2_mean=np.mean(r2_list))
    # 绘制曲线图
    plot_curves(title='Chamber Pressure vs Burned time',
                xlabel='$\it{t}$ /s',  # t 设置为斜体，保留字体
                ylabel='$\it{p}$ /MPa',  # p 设置为斜体，保留字体
                x=np.array(t_array_list).T,
                y=np.array(p_array_list).T * 1e-6,
                ref_x=t_array_target,
                ref_y=p_array_target * 1e-6,
                filename=f"{project_path}/internal_ballistic.svg")

if __name__ == "__main__":
    control_params=get_control_params(question_name="K正弦状燃面装药",
            postfix="_epoch50000",
            cluster_num=5,
            guid_w=0.5)
    start_internal_ballistic_cal(control_params)